#ifndef SM4_H_
#define SM4_H_


#define SM4_RNDS(k0, k1, k2, k3, F)          \
      do {                                   \
         B0 ^= F(B1 ^ B2 ^ B3 ^ ks->rk[k0]); \
         B1 ^= F(B0 ^ B2 ^ B3 ^ ks->rk[k1]); \
         B2 ^= F(B0 ^ B1 ^ B3 ^ ks->rk[k2]); \
         B3 ^= F(B0 ^ B1 ^ B2 ^ ks->rk[k3]); \
      } while(0)

# define SM4_ENCRYPT     1
# define SM4_DECRYPT     0

# define SM4_BLOCK_SIZE    16
# define SM4_KEY_SCHEDULE  32

typedef struct SM4_KEY_st {
    unsigned int rk[SM4_KEY_SCHEDULE];
} SM4_KEY;

namespace A
{

    class SM4
    {
    public:
        int sm4_set_key(const unsigned char *key, SM4_KEY *ks);
        void sm4_encrypt(const unsigned char *in, unsigned char *out, const SM4_KEY *ks);
        void sm4_decrypt(const unsigned char *in, unsigned char *out, const SM4_KEY *ks);

    private:

        //small fun,inline improve speed
        inline unsigned int rotl(unsigned int a, unsigned char n)
        {
            return (a << n) | (a >> (32 - n));
        }

        inline unsigned int load_u32_be(const unsigned char *b, unsigned int n)
        {
            return ((unsigned int)b[4 * n] << 24) |
                   ((unsigned int)b[4 * n + 1] << 16) |
                   ((unsigned int)b[4 * n + 2] << 8) |
                   ((unsigned int)b[4 * n + 3]);
        }

        inline void store_u32_be(unsigned int v, unsigned char *b)
        {
            b[0] = (unsigned char)(v >> 24);
            b[1] = (unsigned char)(v >> 16);
            b[2] = (unsigned char)(v >> 8);
            b[3] = (unsigned char)(v);
        }

        inline int SM4_T_slow(unsigned int X)
        {
            unsigned int t = 0;

            t |= ((unsigned int)SM4_S[(unsigned char)(X >> 24)]) << 24;
            t |= ((unsigned int)SM4_S[(unsigned char)(X >> 16)]) << 16;
            t |= ((unsigned int)SM4_S[(unsigned char)(X >> 8)]) << 8;
            t |= SM4_S[(unsigned char)X];

            /*
             * L linear transform
             */
            return t ^ rotl(t, 2) ^ rotl(t, 10) ^ rotl(t, 18) ^ rotl(t, 24);
        }

        inline unsigned int SM4_T(unsigned int X)
        {
            return SM4_SBOX_T[(unsigned char)(X >> 24)] ^
                   rotl(SM4_SBOX_T[(unsigned char)(X >> 16)], 24) ^
                   rotl(SM4_SBOX_T[(unsigned char)(X >> 8)], 16) ^
                   rotl(SM4_SBOX_T[(unsigned char)X], 8);
        }

    private:
        static const unsigned char SM4_S[256] ;

        /*
         * SM4_SBOX_T[j] == L(SM4_SBOX[j]).
         */
        static const unsigned int SM4_SBOX_T[256];

    };

}
#endif // SM4_H
